// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>

#include <pollux/expression/cast_expr.h>
#include <pollux/expression/coalesce_expr.h>
#include <pollux/expression/conjunct_expr.h>
#include <pollux/expression/constant_expr.h>
#include <pollux/expression/register_special_form.h>
#include <pollux/expression/special_form_registry.h>
#include <pollux/expression/switch_expr.h>
#include <pollux/expression/try_expr.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux;
using namespace kumo::pollux::core;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::test;

class FunctionCallToSpecialFormTest : public testing::Test,
                                      public VectorBuilder {
 protected:
  static void SetUpTestCase() {
    registerFunctionCallToSpecialForms();
    memory::MemoryManager::testingSetInstance({});
  }
  const core::QueryConfig config_{{}};
};

TEST_F(FunctionCallToSpecialFormTest, andCall) {
  ASSERT_TRUE(isFunctionCallToSpecialFormRegistered("and"));

  auto type = resolveTypeForSpecialForm("and", {BOOLEAN(), BOOLEAN()});
  ASSERT_EQ(type, BOOLEAN());

  auto specialForm = constructSpecialForm(
      "and",
      BOOLEAN(),
      {std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<bool>({true})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<bool>({false}))},
      false,
      config_);
  ASSERT_EQ(typeid(*specialForm), typeid(const ConjunctExpr&));
}

TEST_F(FunctionCallToSpecialFormTest, castCall) {
  ASSERT_TRUE(isFunctionCallToSpecialFormRegistered("cast"));

  ASSERT_THROW(resolveTypeForSpecialForm("cast", {}), PolluxRuntimeError);

  auto specialForm = constructSpecialForm(
      "cast",
      DOUBLE(),
      {std::make_shared<ConstantExpr>(
          vectorMaker_.constant_vector<int32_t>({0}))},
      false,
      config_);
  ASSERT_EQ(typeid(*specialForm), typeid(const CastExpr&));
}

TEST_F(FunctionCallToSpecialFormTest, coalesceCall) {
  ASSERT_TRUE(isFunctionCallToSpecialFormRegistered("coalesce"));

  auto type = resolveTypeForSpecialForm("coalesce", {BOOLEAN()});
  ASSERT_EQ(type, BOOLEAN());

  auto specialForm = constructSpecialForm(
      "coalesce",
      INTEGER(),
      {std::make_shared<ConstantExpr>(
          vectorMaker_.constant_vector<int32_t>({0}))},
      false,
      config_);
  ASSERT_EQ(typeid(*specialForm), typeid(const CoalesceExpr&));
}

TEST_F(FunctionCallToSpecialFormTest, ifCall) {
  ASSERT_TRUE(isFunctionCallToSpecialFormRegistered("if"));

  auto type =
      resolveTypeForSpecialForm("if", {BOOLEAN(), INTEGER(), INTEGER()});
  ASSERT_EQ(type, INTEGER());

  auto specialForm = constructSpecialForm(
      "if",
      INTEGER(),
      {std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<bool>({true})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<int32_t>({0})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<int32_t>({1}))},
      false,
      config_);
  ASSERT_EQ(typeid(*specialForm), typeid(const SwitchExpr&));
}

TEST_F(FunctionCallToSpecialFormTest, orCall) {
  ASSERT_TRUE(isFunctionCallToSpecialFormRegistered("or"));

  auto type = resolveTypeForSpecialForm("or", {BOOLEAN(), BOOLEAN()});
  ASSERT_EQ(type, BOOLEAN());

  auto specialForm = constructSpecialForm(
      "or",
      BOOLEAN(),
      {std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<bool>({true})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<bool>({false}))},
      false,
      config_);
  ASSERT_EQ(typeid(*specialForm), typeid(const ConjunctExpr&));
}

TEST_F(FunctionCallToSpecialFormTest, switchCall) {
  ASSERT_TRUE(isFunctionCallToSpecialFormRegistered("switch"));

  auto type = resolveTypeForSpecialForm(
      "switch", {BOOLEAN(), INTEGER(), BOOLEAN(), INTEGER(), INTEGER()});
  ASSERT_EQ(type, INTEGER());

  auto specialForm = constructSpecialForm(
      "switch",
      INTEGER(),
      {std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<bool>({true})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<int32_t>({0})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<bool>({false})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<int32_t>({1})),
       std::make_shared<ConstantExpr>(
           vectorMaker_.constant_vector<int32_t>({2}))},
      false,
      config_);
  ASSERT_EQ(typeid(*specialForm), typeid(const SwitchExpr&));
}

TEST_F(FunctionCallToSpecialFormTest, tryCall) {
  ASSERT_TRUE(isFunctionCallToSpecialFormRegistered("try"));

  auto type = resolveTypeForSpecialForm("try", {BOOLEAN()});
  ASSERT_EQ(type, BOOLEAN());

  auto specialForm = constructSpecialForm(
      "try",
      INTEGER(),
      {std::make_shared<ConstantExpr>(
          vectorMaker_.constant_vector<int32_t>({0}))},
      false,
      config_);
  ASSERT_EQ(typeid(*specialForm), typeid(const TryExpr&));
}

TEST_F(FunctionCallToSpecialFormTest, notASpecialForm) {
  ASSERT_FALSE(isFunctionCallToSpecialFormRegistered("not_a_special_form"));

  auto type = resolveTypeForSpecialForm("not_a_special_form", {BOOLEAN()});
  ASSERT_EQ(type, nullptr);

  auto specialForm = constructSpecialForm(
      "not_a_special_form",
      INTEGER(),
      {std::make_shared<ConstantExpr>(
          vectorMaker_.constant_vector<int32_t>({0}))},
      false,
      config_);
  ASSERT_EQ(specialForm, nullptr);
}

class FunctionCallToSpecialFormSanitizeNameTest : public testing::Test,
                                                  public VectorBuilder {
 protected:
  static void SetUpTestCase() {
    // This class does not pre-register the special forms.
    memory::MemoryManager::testingSetInstance({});
  }
};

TEST_F(FunctionCallToSpecialFormSanitizeNameTest, sanitize_name) {
  // Make sure no special forms are registered.
  unregisterAllFunctionCallToSpecialForm();

  ASSERT_FALSE(isFunctionCallToSpecialFormRegistered("and"));
  ASSERT_FALSE(isFunctionCallToSpecialFormRegistered("AND"));
  ASSERT_FALSE(isFunctionCallToSpecialFormRegistered("or"));
  ASSERT_FALSE(isFunctionCallToSpecialFormRegistered("OR"));

  registerFunctionCallToSpecialForm(
      "and", std::make_unique<ConjunctCallToSpecialForm>(true /* isAnd */));
  registerFunctionCallToSpecialForm(
      "OR", std::make_unique<ConjunctCallToSpecialForm>(false /* isAnd */));

  auto testLookup = [this](const std::string& name) {
    auto type = resolveTypeForSpecialForm(name, {BOOLEAN(), BOOLEAN()});
    ASSERT_EQ(type, BOOLEAN());

    auto specialForm = constructSpecialForm(
        name,
        BOOLEAN(),
        {std::make_shared<ConstantExpr>(
             vectorMaker_.constant_vector<bool>({true})),
         std::make_shared<ConstantExpr>(
             vectorMaker_.constant_vector<bool>({false}))},
        false,
        core::QueryConfig{{}});
    ASSERT_EQ(typeid(*specialForm), typeid(const ConjunctExpr&));
  };

  testLookup("and");
  testLookup("AND");
  testLookup("or");
  testLookup("OR");
}
